{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# High-level Caffe2 Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import caffe2\n",
    "import numpy as np\n",
    "from caffe2.python import core, model_helper, workspace, visualize, brew, optimizer, utils\n",
    "from caffe2.proto import caffe2_pb2\n",
    "from common.params import *\n",
    "from common.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OS:  linux\n",
      "Python:  3.5.2 |Anaconda custom (64-bit)| (default, Jul  2 2016, 17:53:06) \n",
      "[GCC 4.4.7 20120313 (Red Hat 4.4.7-1)]\n",
      "Numpy:  1.13.3\n",
      "GPU:  ['Tesla K80']\n"
     ]
    }
   ],
   "source": [
    "print(\"OS: \", sys.platform)\n",
    "print(\"Python: \", sys.version)\n",
    "print(\"Numpy: \", np.__version__)\n",
    "print(\"GPU: \", get_gpu_name())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if GPU:\n",
    "    device_opts = core.DeviceOption(caffe2_pb2.CUDA, 0)  # Run on GPU\n",
    "else:\n",
    "    device_opts = core.DeviceOption(caffe2_pb2.CPU, 0)  # Run on CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model(m, device_opts) :\n",
    "    with core.DeviceScope(device_opts):\n",
    "        conv1 = brew.conv(m, 'data', 'conv1', dim_in=3, dim_out=50, kernel=3, pad=1, no_gradient_to_input=1)\n",
    "        relu1 = brew.relu(m, conv1, 'relu1')\n",
    "        conv2 = brew.conv(m, relu1, 'conv2', dim_in=50, dim_out=50, kernel=3, pad=1)\n",
    "        pool1 = brew.max_pool(m, conv2, 'pool1', kernel=2, stride=2)\n",
    "        relu2 = brew.relu(m, pool1, 'relu2')\n",
    "        drop1 = brew.dropout(m, relu2, 'drop1', ratio=0.25)\n",
    "\n",
    "        conv3 = brew.conv(m, drop1, 'conv3', dim_in=50, dim_out=100, kernel=3, pad=1)\n",
    "        relu3 = brew.relu(m, conv3, 'relu3')\n",
    "        conv4 = brew.conv(m, relu3, 'conv4', dim_in=100, dim_out=100, kernel=3, pad=1)\n",
    "        pool2 = brew.max_pool(m, conv4, 'pool2', kernel=2, stride=2)   \n",
    "        relu4 = brew.relu(m, pool2, 'relu4')\n",
    "        drop2 = brew.dropout(m, relu4, 'drop2', ratio=0.25)\n",
    "        \n",
    "        fc1 = brew.fc(m, drop2, 'fc1', dim_in=100 * 8 * 8, dim_out=512)\n",
    "        relu5 = brew.relu(m, fc1, 'relu5')\n",
    "        drop3 = brew.dropout(m, relu5, 'drop3', ratio=0.5)\n",
    "        \n",
    "        fc2 = brew.fc(m, drop3, 'fc2', dim_in=512, dim_out=N_CLASSES)\n",
    "        softmax = brew.softmax(m, fc2, 'softmax')\n",
    "        return softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_training_operators(softmax, m, device_opts) :\n",
    "    with core.DeviceScope(device_opts):\n",
    "        xent = m.LabelCrossEntropy([softmax, \"label\"], 'xent')\n",
    "        loss = m.AveragedLoss(xent, \"loss\")\n",
    "        #brew.accuracy(m, [softmax, \"label\"], \"accuracy\")\n",
    "        m.AddGradientOperators([loss])\n",
    "        opt = optimizer.build_sgd(\n",
    "            m,\n",
    "            base_learning_rate=LR, \n",
    "            policy='fixed',\n",
    "            momentum=MOMENTUM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_model():\n",
    "    # Create Place-holder for data\n",
    "    workspace.FeedBlob(\"data\", x_train[:BATCHSIZE], device_option=device_opts)\n",
    "    workspace.FeedBlob(\"label\", y_train[:BATCHSIZE], device_option=device_opts)\n",
    "    \n",
    "    # Initialise model\n",
    "    train_arg_scope = {\n",
    "        'order': 'NCHW',\n",
    "        'use_cudnn': True,\n",
    "        'cudnn_exhaustive_search': True,\n",
    "        'ws_nbytes_limit': (64 * 1024 * 1024),\n",
    "    }\n",
    "    train_model = model_helper.ModelHelper(\n",
    "        name=\"train_net\", arg_scope=train_arg_scope\n",
    "    )\n",
    "    softmax = create_model(train_model, device_opts=device_opts)\n",
    "    add_training_operators(softmax, train_model, device_opts=device_opts)\n",
    "\n",
    "    # Initialise workspace\n",
    "    workspace.RunNetOnce(train_model.param_init_net)\n",
    "    workspace.CreateNet(train_model.net)\n",
    "    return train_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing train set...\n",
      "Preparing test set...\n",
      "(50000, 3, 32, 32) (10000, 3, 32, 32) (50000,) (10000,)\n",
      "float32 float32 int32 int32\n",
      "CPU times: user 836 ms, sys: 550 ms, total: 1.39 s\n",
      "Wall time: 1.38 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Data into format for library\n",
    "x_train, x_test, y_train, y_test = cifar_for_library(channel_first=True)\n",
    "print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)\n",
    "print(x_train.dtype, x_test.dtype, y_train.dtype, y_test.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 424 ms, sys: 330 ms, total: 754 ms\n",
      "Wall time: 762 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Initialise model\n",
    "model = init_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:  0\n",
      "0: 1.4412875175476074\n",
      "Finished epoch:  1\n",
      "1: 1.0683350563049316\n",
      "Finished epoch:  2\n",
      "2: 0.9755038619041443\n",
      "Finished epoch:  3\n",
      "3: 0.7842862010002136\n",
      "Finished epoch:  4\n",
      "4: 0.7359310388565063\n",
      "Finished epoch:  5\n",
      "5: 0.7806798815727234\n",
      "Finished epoch:  6\n",
      "6: 0.6619545221328735\n",
      "Finished epoch:  7\n",
      "7: 0.56577068567276\n",
      "Finished epoch:  8\n",
      "8: 0.5507668256759644\n",
      "Finished epoch:  9\n",
      "9: 0.44708868861198425\n",
      "CPU times: user 2min 3s, sys: 25.1 s, total: 2min 28s\n",
      "Wall time: 2min 28s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# 148s\n",
    "# Train model\n",
    "for j in range(EPOCHS):\n",
    "    for data, label in yield_mb(x_train, y_train, BATCHSIZE, shuffle=True):\n",
    "        # Run one mini-batch at time\n",
    "        workspace.FeedBlob(\"data\", data, device_option=device_opts)\n",
    "        workspace.FeedBlob(\"label\", label, device_option=device_opts)\n",
    "        workspace.RunNet(model.net)       \n",
    "    print(\"Finished epoch: \", j)\n",
    "    print(str(j) + ': ' + str(workspace.FetchBlob(\"loss\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 867 ms, sys: 217 ms, total: 1.08 s\n",
      "Wall time: 1.15 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Init test model\n",
    "test_arg_scope = {\n",
    "    'order': 'NCHW',\n",
    "    'use_cudnn': True,\n",
    "    'cudnn_exhaustive_search': True,\n",
    "    'ws_nbytes_limit': (64 * 1024 * 1024),\n",
    "    'is_test': True,\n",
    "}\n",
    "test_model= model_helper.ModelHelper(name=\"test_net\", init_params=False, arg_scope=test_arg_scope)\n",
    "create_model(test_model, device_opts=device_opts)\n",
    "workspace.RunNetOnce(test_model.param_init_net)\n",
    "workspace.CreateNet(test_model.net, overwrite=True)\n",
    "\n",
    "# Run test\n",
    "n_samples = (y_test.shape[0]//BATCHSIZE)*BATCHSIZE\n",
    "y_guess = np.zeros(n_samples, dtype=np.int)\n",
    "y_truth = y_test[:n_samples]\n",
    "c = 0\n",
    "for data, label in yield_mb(x_test, y_test, BATCHSIZE):\n",
    "    workspace.FeedBlob(\"data\", data, device_option=device_opts)\n",
    "    workspace.RunNet(test_model.net)\n",
    "    y_guess[c*BATCHSIZE:(c+1)*BATCHSIZE] = (np.argmax(workspace.FetchBlob(\"softmax\"), axis=-1))\n",
    "    c += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy:  0.790564903846\n"
     ]
    }
   ],
   "source": [
    "print(\"Accuracy: \", sum(y_guess == y_truth)/float(len(y_guess)))"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:py35]",
   "language": "python",
   "name": "conda-env-py35-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
